module Skit
  module Health
    class Server
      attr_accessor :connections, :check

      def initialize(port_or_socket, &block)
        @port, @socket = extract_port_and_socket(port_or_socket)
        @host = '127.0.0.1'
        @check = block
        @connections = []
        @stopped = false

        debug("Starting health check server on #{server_info}")
        debug("Initial run: #{Skit::Health.run_check(block)}")
      end

      def start
        require_em!
        @signature = EventMachine.start_server(*server_args) do |con|
          debug(" [+] -> new connection #{con.object_id}")
          @connections << con
          con.server = self
        end
      end

      def stop
        if @connections.empty?
          unless @stopped
            debug('No connections left, shutting down ...')
            EventMachine.stop_server(@signature)
            EventMachine.stop
            @stopped = true
          end
        else
          debug("Waiting for #{@connections.size} connection(s) to finish ...")
          EventMachine.add_timer(1) { stop }
        end
      end


      def debug(message)
        Health.debug(message)
      end

      private

      def extract_port_and_socket(port_or_socket)
        port_or_socket.is_a?(Integer) ?
          [port_or_socket, nil] :
          [nil, port_or_socket]
      end

      def server_info
         @port ? "#{@host}:#{@port}" : @socket
      end

      def server_args
        @port ? [@host, @port, HttpConnection] : [@socket, UnixConnection]
      end

      def require_em!
        begin
          require 'eventmachine'
        rescue LoadError
          $stderr.puts "You need eventmachine ~> 1.0 gem " +
            "to run Skit::Health::Server. " +
            "Please add it to Gemfile and run 'bundle install'."
          exit 1
        end
      end
    end

    class BaseConnection < EventMachine::Connection
      CR = "\n"

      attr_accessor :server

      def check_result
        "#{Skit::Health.run_check(server.check)}#{CR}"
      end

      def send(&block)
        response = check_result
        yield response
        server.debug("#{self.object_id} response: #{response}")
        close_connection_after_writing
      end

      def unbind
        server.debug(" [-] -> removing connection #{self.object_id}")
        server.connections.delete(self)
      end
    end

    class HttpConnection < BaseConnection
      def receive_data(data)
        send do |response|
          send_data headers(response).join(CR)
          send_data CR * 2
          send_data response
        end
      end

      def headers(response)
        [
          'HTTP/1.1 200 OK',
          'Content-Type: text/plain; charset=UTF-8',
          'Connection: close',
          "Content-Length: #{response.length}",
        ]
      end
    end

    class UnixConnection < BaseConnection
      def receive_data(data)
        send do |response|
          send_data response
        end
      end
    end
  end
end
